import os
import shutil
from pathlib import Path
import io
import imghdr
import Logger_config
from docx import Document
from docx.enum.text import WD_ALIGN_PARAGRAPH, WD_COLOR_INDEX
from docx.shared import Pt, RGBColor
from docx.oxml.shared import qn
from Email import Email
from ImageUtil import ImageUtil

logger = Logger_config.setup_logger()


class WordContentExtractor:
    """Word文档内容提取器"""

    def __init__(self):
        """初始化"""
        self.temp_dir = os.path.join(os.path.dirname(__file__), 'temp')
        os.makedirs(self.temp_dir, exist_ok=True)

    def extract_content(self, docx_path):
        """
        从Word文档中提取文本并构建HTML正文，同时收集CID图片信息用于邮件附件

        参数:
            docx_path: Word文档路径

        返回:
            tuple: (HTML格式文本内容, image_cid_list)
                   image_cid_list: [{'path': 路径, 'cid': CID, 'mime_type': MIME类型}, ...]
        """
        try:
            doc = Document(docx_path)
            text_content = ""
            image_cid_list = []

            # 定义 CSS 样式模板
            css_style = """
                <style>
                    body {
                        line-height: 1.6;
                        background-color: #ffffff;
                        margin: 0;
                        padding: 20px;
                    }
                    p {
                        margin: 0 0 10px 0;
                        padding: 0;
                    }
                    img {
                        max-width: 20%;
                        height: auto;
                        display: block;
                        margin: 10px auto;
                    }
                </style>
            """

            html_header = f"""
                <!DOCTYPE html>
                <html>
                <head>
                    <meta charset="UTF-8">
                    <title>邮件正文</title>
                    {css_style}
                </head>
                <body>
            """

            html_footer = """
                </body>
                </html>
            """

            for para in doc.paragraphs:
                # 获取段落样式
                para_format = para.paragraph_format
                para_style = {
                    'font_name': para.style.font.name,
                    'font_size': para.style.font.size.pt if para.style.font.size else 12,
                    'color': para.style.font.color.rgb if para.style.font.color else '#000000',
                    'alignment': para.alignment,

                    # 缩进
                    'left_indent': para_format.left_indent.pt if para_format.left_indent else 0,
                    'right_indent': para_format.right_indent.pt if para_format.right_indent else 0
                }

                align_map = {
                    None: 'left',
                    WD_ALIGN_PARAGRAPH.LEFT: 'left',
                    WD_ALIGN_PARAGRAPH.CENTER: 'center',
                    WD_ALIGN_PARAGRAPH.RIGHT: 'right',
                    WD_ALIGN_PARAGRAPH.JUSTIFY: 'justify'
                }

                p_style_attrs = []
                if para_style['font_name']:
                    p_style_attrs.append(f'font-family:{para_style["font_name"]}')
                if para_style['font_size']:
                    p_style_attrs.append(f'font-size:{para_style["font_size"]}pt')
                if para_style['color']:
                    p_style_attrs.append(f'color:{para_style["color"]}')
                if para_style['alignment'] in align_map:
                    p_style_attrs.append(f'text-align:{align_map[para_style["alignment"]]}')

                text_content += f'<p style="{";".join(p_style_attrs)}">'

                for run in para.runs:
                    run_style = {
                        'bold': run.bold,
                        'italic': run.italic,
                        'underline': run.underline,
                        'font_name': run.font.name,
                        'font_size': run.font.size.pt if run.font.size else para_style['font_size'],
                        'color': run.font.color.rgb if run.font.color else para_style['color'],
                        'highlight_color': run.font.highlight_color,
                    }

                    span_styles = []
                    if run_style['font_name']:
                        span_styles.append(f'font-family:{run_style["font_name"]}')
                    if run_style['font_size']:
                        span_styles.append(f'font-size:{run_style["font_size"]}pt')
                    if run_style['color']:
                        span_styles.append(f'color:{run_style["color"]}')
                    if run_style['bold']:
                        span_styles.append('font-weight:bold')
                    if run_style['italic']:
                        span_styles.append('font-style:italic')
                    if run_style['underline']:
                        span_styles.append('text-decoration:underline')
                    if run_style['highlight_color']:
                        highlight_color = self._map_highlight_color(run_style['highlight_color'])
                        if highlight_color:
                            span_styles.append(f'background-color:{highlight_color} !important')

                    text_content += f'<span style="{";".join(span_styles)}">'
                    text_content += run.text
                    text_content += '</span>'

                    # 提取图片（现在返回的是 bytes）
                    image_data = self._extract_image_from_shape(run._element, doc.part.rels)
                    if image_data:
                        # 检测格式
                        img_format = self._detect_image_format(image_data)
                        mime_type = f"image/{img_format}" if img_format else "image/png"
                        cid = f"image{len(image_cid_list)}"

                        # 调用通用压缩方法，并传入目标大小（例如 512KB）
                        compressed_data, _ = ImageUtil._compress_image(image_data, max_size_kb=300)

                        image_cid_list.append({
                            'data': compressed_data,
                            'cid': cid,
                            'mime_type': mime_type
                        })

                        text_content += f'<img src="cid:{cid}" style="max-width:30%; height:auto; display:block;margin:10px auto;">'

                text_content += '</p>\n'

            full_html = html_header + f"<div>{text_content}</div>" + html_footer
            logger.info(f"成功从 {docx_path} 提取内容")
            return full_html, image_cid_list

        except Exception as e:
            logger.error(f"提取内容失败: {str(e)}")
            raise

    def _map_highlight_color(self, color):
        """将 docx 的 HighlightColor 映射为 CSS 颜色值"""
        color_map = {
            WD_COLOR_INDEX.YELLOW: '#FFFF00',
            WD_COLOR_INDEX.GREEN: '#00FF00',
            WD_COLOR_INDEX.CYAN: '#00FFFF',
            WD_COLOR_INDEX.MAGENTA: '#FF00FF',
            WD_COLOR_INDEX.BLUE: '#0000FF',
            WD_COLOR_INDEX.RED: '#FF0000',
            WD_COLOR_INDEX.DARK_BLUE: '#00008B',
            WD_COLOR_INDEX.DARK_RED: '#8B0000',
            WD_COLOR_INDEX.DARK_YELLOW: '#CCCC00',
            WD_COLOR_INDEX.GRAY_25: '#C0C0C0',
            WD_COLOR_INDEX.GRAY_50: '#808080',
        }
        if color in color_map:
            return color_map[color]
        elif isinstance(color, RGBColor):
            return f'#{color.hex}'
        return None

    def _extract_image_from_shape(self, run_element, rels):
        """
        从 run 元素中提取图片并返回二进制数据

        参数:
            run_element: run 的 XML 元素
            rels: 文档关系对象

        返回:
            bytes: 图像的原始二进制数据
        """
        try:
            ns = {
                'a': 'http://schemas.openxmlformats.org/drawingml/2006/main',
                'r': 'http://schemas.openxmlformats.org/officeDocument/2006/relationships',
                'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main',
                'v': 'urn:schemas-microsoft-com:vml',
                'pic': 'http://schemas.openxmlformats.org/drawingml/2006/presentationml'
            }

            # 查找现代图形格式 <a:blip r:embed="..."/>
            blips = run_element.findall(f'.//{{{ns["a"]}}}blip')
            for blip in blips:
                embed_rel_id = blip.get(f'{{{ns["r"]}}}embed')
                if embed_rel_id and embed_rel_id in rels:
                    image_part = rels[embed_rel_id].target_part
                    image_data = image_part.blob
                    return image_data

            # 查找 VML 图形格式 <v:imagedata r:id="..."/>
            v_images = run_element.findall(f'.//{{{ns["v"]}}}imagedata')
            for vimg in v_images:
                rel_id = vimg.get(f'{{{ns["r"]}}}id')
                if rel_id and rel_id in rels:
                    image_part = rels[rel_id].target_part
                    image_data = image_part.blob
                    return image_data

            return None
        except Exception as e:
            logger.error(f"图片提取失败: {str(e)}", exc_info=True)
            return None

    def _detect_image_format(self, img_data):
        """
        检测图片数据的原始格式

        参数:
            img_data: 图片的原始二进制数据

        返回:
            str: 图片格式(png/jpg/jpeg/gif等)，无法识别时返回None
        """
        try:
            img_format = imghdr.what(None, h=img_data)
            if img_format == 'jpeg':
                return 'jpg'
            return img_format
        except Exception as e:
            logger.warning(f"图片格式检测失败: {str(e)}")
            return None

    def cleanup(self):
        """清理临时目录"""
        try:
            if os.path.exists(self.temp_dir):
                shutil.rmtree(self.temp_dir)
        except Exception as e:
            logger.warning(f"清理临时目录失败: {str(e)}")


if __name__ == "__main__":
    extractor = WordContentExtractor()

    try:
        input_docx = r"d:\cmsr\10_上研院\CodingSpace\Python\ai-docs\data\input\部门简讯\产业数智化产品部部门简讯（第18-19周，4月28日-5月9日）-邮件正文.docx"
        html_content, images = extractor.extract_content(input_docx)

        logger.info(f"提取的HTML内容长度: {html_content}")
        logger.info(f"提取的图片数量: {len(images)}")

        subject = "测试邮件 - 2025年产业数智化产品部部门简讯"
        recipients = ["chengyuansen@cmsr.chinamobile.com"]
        cc_recipients = []
        email_sender = Email()
        email_sender.send_email_with_attachments(subject, recipients, cc_recipients, html_content, images)
        print("\n测试邮件已发送")

    finally:
        extractor.cleanup()